# 在html页面的温度滑块旁边添加一个下拉菜单，选中ollama模型的名称model_name，包括deepseek-r1:1.5b，qwen2.5:3b等，
# 选中后同步修改app.py中的变量model_name。只显示改动的代码


from flask import Flask, render_template, request, jsonify
import requests
import sqlite3
from datetime import datetime
import os
import csv

app = Flask(__name__)


# Ollama的 API地址和密钥
API_URL = "http://localhost:11434/v1/chat/completions"
API_KEY = "ollama"
model_name = "llama3.2"  


# 全局变量存储当前数据库路径
dir_path = os.path.dirname(os.path.realpath(__file__))
current_db_path = os.path.join(dir_path, 'chat_history.db')
# current_db_name = 'chat_history.db'

# 调用chat_models 模式
def call_zhipuai(messages, temperature, model_name):
    headers = {
        "Authorization": f"Bearer {API_KEY}",
        "Content-Type": "application/json",
    }
    data = {
        "model": model_name,  # 使用的模型
        "messages": messages,      # 对话上下文
        "temperature": temperature,  # 温度参数
        "stream": False            # 非流式返回
    }
    response = requests.post(API_URL, headers=headers, json=data)
    if response.status_code == 200:
        return response.json()["choices"][0]["message"]["content"]
    else:
        return f"Error: {response.status_code} - {response.text}"

# 保存对话历史到 SQLite 数据库
def save_chat_history(db_path, user_input, ai_response):
    conn = sqlite3.connect(db_path)
    c = conn.cursor()
    c.execute('''CREATE TABLE IF NOT EXISTS chat_history
                 (id INTEGER PRIMARY KEY AUTOINCREMENT,
                  timestamp TEXT,
                  user_input TEXT,
                  ai_response TEXT)''')
    timestamp = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
    c.execute("INSERT INTO chat_history (timestamp, user_input, ai_response) VALUES (?, ?, ?)",
              (timestamp, user_input, ai_response))
    conn.commit()
    conn.close()

# 加载对话历史从 SQLite 数据库
def load_chat_history(db_path):
    if not os.path.exists(db_path):
        return []
    conn = sqlite3.connect(db_path)
    c = conn.cursor()
    c.execute("SELECT timestamp, user_input, ai_response FROM chat_history")
    history = c.fetchall()
    conn.close()
    return history

# 导出历史记录为 CSV 文件
def export_to_csv(db_path, csv_path):
    if not os.path.exists(db_path):
        return []

    history = load_chat_history(db_path)
    
    with open(csv_path, mode='w', newline='', encoding='utf-8') as file:
        writer = csv.writer(file)
        writer.writerow(['Timestamp', 'User Input', 'AI Response'])
        for record in history:
            writer.writerow(record)

# 主页路由
@app.route('/')
def index():
    return render_template('index_app04.html')

# 聊天路由
@app.route('/chat', methods=['POST'])
def chat():
    global current_db_path

    user_input = request.json.get('message')
    temperature = float(request.json.get('temperature', 0.7))
    model_name = request.json.get('model_name', "qwen2.5:3b")
    session_messages = request.json.get('messages', [])
    session_messages.append({"role": "user", "content": user_input})

    # 调用智谱AI
    ai_response = call_zhipuai(session_messages, temperature, model_name)
    session_messages.append({"role": "assistant", "content": ai_response})

    # 如果当前数据库路径存在，保存对话历史
    if current_db_path:
        save_chat_history(current_db_path, user_input, ai_response)

    return jsonify({"response": ai_response, "messages": session_messages})

# 保存数据库路由
@app.route('/save_db', methods=['POST'])
def save_db():
    global current_db_path
    db_path = request.json.get('db_path')
    if not db_path:
        # 默认数据库文件名
        # db_path = f"chat_{datetime.now().strftime('%Y%m%d_%H%M%S')}.db"
        db_path = current_db_path
    else:
        db_path = os.path.join(dir_path, db_path) 
    
    if current_db_path != db_path:
        import shutil
        shutil.copy(current_db_path, db_path)
    
    current_db_path = db_path
    
    return jsonify({"status": "success", "db_path": db_path})

# 加载数据库路由
@app.route('/load_db', methods=['POST'])
def load_db():
    global current_db_path
    db_path = request.json.get('db_path')
    if not db_path:
        db_path = current_db_path
    else:
        db_path = os.path.join(dir_path, db_path) 
    if not os.path.exists(db_path):
        return jsonify({"status": "error", "message": "数据库文件不存在"})
    current_db_path = db_path
    history = load_chat_history(db_path)
    return jsonify({"status": "success", "history": history})

# 导出 CSV 路由
@app.route('/export_csv', methods=['POST'])
def export_csv():
    global current_db_path
    csv_path = request.json.get('csv_path')
    if not csv_path:
  
        csv_path = f"chat_{datetime.now().strftime('%Y%m%d_%H%M')}.csv"
        csv_path = os.path.join(dir_path, csv_path) 
    else: 
        csv_path = os.path.join(dir_path, csv_path) 
   
    try:
        export_to_csv(current_db_path, csv_path)
        return jsonify({"status": "success", "csv_path": csv_path})
    except Exception as e:
        return jsonify({"status": "error", "message": str(e)})

if __name__ == '__main__':
    app.run(debug=True)